import { app } from "../../scripts/app.js";

/**
 * 通用训练节点UI注册函数
 * @param {Object} options
 *   comfyClass: 节点类型
 *   eventName: 事件名（如"flux_training_progress"）
 *   themeColor: 主题色（如#8b5cf6）
 *   chartTitle: 图表标题
 *   progressBarColor: 进度条颜色
 *   gradientColor: 图表渐变色
 *   legendIcon: 图例标识（如"FLUX"、"SDXL"等）
 *   titleFont: 标题字体
 *   lossLineColor: loss曲线颜色
 */
export function registerTrainingNodeUI(options) {
    const {
        comfyClass,
        eventName,
        themeColor = "#8b5cf6",
        chartTitle = "Training Loss",
        progressBarColor = "#8b5cf6",
        gradientColor = "#8b5cf660",
        legendIcon = "FLUX",
        titleFont = "bold 16px Arial",
        lossLineColor = null
    } = options;

    console.log(`[registerTrainingNodeUI] 注册UI: comfyClass=${comfyClass}, eventName=${eventName}`);

    app.registerExtension({
        name: `Comfy.${comfyClass}TrainingUICommon`,
        async beforeRegisterNodeDef(nodeType) {
            if (nodeType.comfyClass !== comfyClass) return;
            console.log(`[registerTrainingNodeUI] beforeRegisterNodeDef: ${comfyClass}`);
            // 统一节点大小
            nodeType.size = [240, 280];

            // onNodeCreated
            const onNodeCreated = nodeType.prototype.onNodeCreated;
            nodeType.prototype.onNodeCreated = function() {
                onNodeCreated?.apply(this, arguments);
                this.lossData = [];
                this.maxDataPoints = 100;
                this.chartVisible = true;
                this.legendHitboxes = {};
            };

            // onDrawForeground
            const onDrawForeground = nodeType.prototype.onDrawForeground;
            nodeType.prototype.onDrawForeground = function(ctx) {
                // 清空画布，防止内容叠加
                ctx.clearRect(0, 0, this.size[0], this.size[1]);
                if (onDrawForeground) onDrawForeground.call(this, ctx);
                if (!this.progress) return;
                const { current, total, status, loss } = this.progress;
                const progress = total > 0 ? current / total : 0;

                // 进度条
                const barHeight = 10;
                const barY = this.size[1] - barHeight - 4;
                ctx.fillStyle = "#2a2a2a";
                ctx.fillRect(0, barY, this.size[0], barHeight);
                // 状态色：缓存中为橙色，训练中为主题色
                const statusColor = status === "缓存中" ? "#f59e42" : progressBarColor;
                ctx.fillStyle = statusColor;
                ctx.fillRect(0, barY, this.size[0] * progress, barHeight);

                // 状态文本
                ctx.fillStyle = status === "缓存中" ? "#f59e42" : themeColor;
                ctx.font = "14px Arial";
                ctx.textAlign = "left";
                ctx.fillText(status || "训练中", 5, barY - 8);
                ctx.textAlign = "right";
                ctx.fillStyle = "#fff";
                ctx.fillText(`${current}/${total}`, this.size[0] - 5, barY - 8);

                // Loss文本（只有当loss有效且不是-1时才显示）
                if (loss !== undefined && loss !== null && loss !== -1 && loss >= 0) {
                    ctx.textAlign = "center";
                    ctx.fillStyle = themeColor;
                    ctx.fillText(`Loss: ${loss.toFixed(4)}`, this.size[0] / 2, barY + barHeight + 16);
                }

                // 只在这里调用一次图表绘制（只有当有有效的loss数据时才绘制）
                if (this.chartVisible && this.lossData.length > 0) {
                    // 检查是否有有效的loss数据（不是-1）
                    const validLossData = this.lossData.filter(d => d.loss >= 0);
                    if (validLossData.length > 0) {
                        this.drawLossChart(ctx);
                        // 只有当有有效数据时才绘制图例
                        this.drawChartLegend(ctx, { x: 30, y: 40, width: this.size[0] - 40, height: this.size[1] - 80 });
                    }
                }
            };

            // onMouseDown
            const onMouseDown = nodeType.prototype.onMouseDown;
            nodeType.prototype.onMouseDown = function(e, pos) {
                if (onMouseDown) onMouseDown.call(this, e, pos);
                for (const id in this.legendHitboxes) {
                    const box = this.legendHitboxes[id];
                    if (pos[0] > box.x && pos[0] < box.x + box.width &&
                        pos[1] > box.y && pos[1] < box.y + box.height) {
                        if (id === "chart") {
                            this.chartVisible = !this.chartVisible;
                        }
                        this.setDirtyCanvas(true, false);
                        return;
                    }
                }
            };

            // updateProgress
            nodeType.prototype.updateProgress = function(data) {
                // 检查是否是新一轮训练，清空lossData
                if ((data.current_step === 1 || data.current_step === 0) && this.lossData.length > 0) {
                    this.lossData = [];
                }
                this.progress = {
                    current: data.current_step || 0,
                    total: data.total_steps || 0,
                    status: data.status || "训练中",
                    loss: data.loss
                };
                // 只有在loss有效且不是-1时才push lossData（-1表示非训练阶段）
                if (data.loss !== undefined && data.loss !== null && data.loss !== -1 && data.loss >= 0) {
                    this.lossData.push({
                        step: data.current_step || 0,
                        loss: data.loss,
                        timestamp: Date.now()
                    });
                    if (this.lossData.length > this.maxDataPoints) {
                        this.lossData = this.lossData.slice(-this.maxDataPoints);
                    }
                }
                this.setDirtyCanvas(true, true);
            };

            // Loss图表相关
            nodeType.prototype.drawLossChart = function(ctx) {
                ctx.clearRect(0, 0, this.size[0], this.size[1]);
                
                // 过滤出有效的loss数据（不是-1）
                const validLossData = this.lossData.filter(d => d.loss >= 0);
                if (validLossData.length < 2) return;
                
                const losses = validLossData.map(d => d.loss);
                const minLoss = Math.min(...losses);
                const maxLoss = Math.max(...losses);
                const lossRange = maxLoss - minLoss || 1e-6;
                const width = this.size[0] - 40;
                const height = this.size[1] - 80;
                const rect = { x: 30, y: 40, width, height };
                const xStep = validLossData.length > 1 ? width / (validLossData.length - 1) : width;

                ctx.save();

                // 坐标轴
                ctx.strokeStyle = "#444";
                ctx.lineWidth = 1;
                ctx.beginPath();
                ctx.moveTo(rect.x, rect.y);
                ctx.lineTo(rect.x, rect.y + height);
                ctx.lineTo(rect.x + width, rect.y + height);
                ctx.stroke();

                // 渐变填充
                const grad = ctx.createLinearGradient(0, rect.y, 0, rect.y + height);
                grad.addColorStop(0, "#3b82f660");
                grad.addColorStop(1, "#3b82f600");
                ctx.fillStyle = grad;
                ctx.beginPath();
                validLossData.forEach((d, i) => {
                    const x = rect.x + i * xStep;
                    const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
                    if (i === 0) ctx.moveTo(x, y);
                    else ctx.lineTo(x, y);
                });
                ctx.lineTo(rect.x + width, rect.y + height);
                ctx.lineTo(rect.x, rect.y + height);
                ctx.closePath();
                ctx.fill();

                // 主线
                ctx.strokeStyle = "#3b82f6";
                ctx.lineWidth = 2;
                ctx.beginPath();
                validLossData.forEach((d, i) => {
                    const x = rect.x + i * xStep;
                    const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
                    if (i === 0) ctx.moveTo(x, y);
                    else ctx.lineTo(x, y);
                });
                ctx.stroke();

                // 数据点
                ctx.fillStyle = "#3b82f6";
                validLossData.forEach((d, i) => {
                    const x = rect.x + i * xStep;
                    const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
                    ctx.beginPath();
                    ctx.arc(x, y, 2, 0, 2 * Math.PI);
                    ctx.fill();
                });

                ctx.restore();
            };
            nodeType.prototype.drawChartTitle = function(ctx) {
                ctx.fillStyle = "#fff";
                ctx.font = titleFont;
                ctx.textAlign = "center";
                let title = chartTitle + " (训练中)";
                if (this.progress && this.progress.status) {
                    title = `${chartTitle} (${this.progress.status})`;
                }
                ctx.fillText(title, this.size[0] / 2, 62);
            };
            nodeType.prototype.drawChartGrid = function(ctx, rect) {
                ctx.strokeStyle = "#333";
                ctx.lineWidth = 1;
                ctx.font = "10px Arial";
                ctx.fillStyle = "#888";
                ctx.textAlign = "right";
                
                // 过滤出有效的loss数据
                const validLossData = this.lossData.filter(d => d.loss >= 0);
                if (validLossData.length < 2) return;
                
                const losses = validLossData.map(d => d.loss);
                const minLoss = Math.min(...losses);
                const maxLoss = Math.max(...losses);
                const lossRange = maxLoss - minLoss;
                for (let i = 0; i <= 4; i++) {
                    const p = i / 4;
                    const y = rect.y + p * rect.height;
                    const value = maxLoss - p * lossRange;
                    ctx.beginPath();
                    ctx.moveTo(rect.x, y);
                    ctx.lineTo(rect.x + rect.width, y);
                    ctx.stroke();
                    ctx.fillText(`${value.toFixed(4)}`, rect.x - 5, y + 3);
                }
                const steps = validLossData.map(d => d.step);
                const minStep = Math.min(...steps);
                const maxStep = Math.max(...steps);
                const stepRange = maxStep - minStep;
                for (let i = 0; i <= 4; i++) {
                    const p = i / 4;
                    const x = rect.x + p * rect.width;
                    const value = minStep + p * stepRange;
                    ctx.beginPath();
                    ctx.moveTo(x, rect.y);
                    ctx.lineTo(x, rect.y + rect.height);
                    ctx.stroke();
                    ctx.textAlign = "center";
                    ctx.fillText(`${Math.round(value)}`, x, rect.y + rect.height + 16);
                }
                ctx.textAlign = "left";
            };
            nodeType.prototype.drawLossLine = function(ctx, rect) {
                // 过滤出有效的loss数据
                const validLossData = this.lossData.filter(d => d.loss >= 0);
                if (validLossData.length < 2) return;
                
                const losses = validLossData.map(d => d.loss);
                const steps = validLossData.map(d => d.step);
                const minLoss = Math.min(...losses);
                const maxLoss = Math.max(...losses);
                const lossRange = maxLoss - minLoss;
                const minStep = Math.min(...steps);
                const maxStep = Math.max(...steps);
                const width = rect.width;
                const height = rect.height;
                const xStep = validLossData.length > 1 ? width / (validLossData.length - 1) : width;
                // 渐变填充
                const grad = ctx.createLinearGradient(0, rect.y, 0, rect.y + height);
                grad.addColorStop(0, gradientColor);
                grad.addColorStop(1, themeColor + "00");
                ctx.fillStyle = grad;
                ctx.beginPath();
                ctx.moveTo(rect.x, rect.y + height);
                validLossData.forEach((d, i) => {
                    const x = rect.x + i * xStep;
                    const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
                    ctx.lineTo(x, y);
                });
                ctx.lineTo(rect.x + (validLossData.length-1) * xStep, rect.y + height);
                ctx.closePath();
                ctx.fill();
                // 曲线
                ctx.strokeStyle = lossLineColor || themeColor;
                ctx.lineWidth = 2;
                ctx.beginPath();
                validLossData.forEach((d, i) => {
                    const x = rect.x + i * xStep;
                    const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
                    if (i === 0) ctx.moveTo(x, y);
                    else ctx.lineTo(x, y);
                });
                ctx.stroke();
                // 数据点
                ctx.fillStyle = themeColor;
                validLossData.forEach((d, i) => {
                    const x = rect.x + i * xStep;
                    const y = rect.y + height - ((d.loss - minLoss) / lossRange) * height;
                    ctx.beginPath();
                    ctx.arc(x, y, 2, 0, 2 * Math.PI);
                    ctx.fill();
                });
            };
            nodeType.prototype.drawChartLegend = function(ctx, rect) {
                const legendY = rect.y + rect.height + 35;
                ctx.font = "12px Arial";
                let currentX = rect.x;
                this.legendHitboxes = {};
                const isVisible = this.chartVisible;
                ctx.globalAlpha = isVisible ? 1.0 : 0.4;
                
                // 过滤出有效的loss数据
                const validLossData = this.lossData.filter(d => d.loss >= 0);
                
                // 只有当有有效的loss数据时才显示loss信息
                if (validLossData.length > 0) {
                    const currentLoss = validLossData[validLossData.length - 1].loss;
                    ctx.fillStyle = themeColor;
                    ctx.fillRect(currentX, legendY - 8, 10, 10);
                    ctx.fillStyle = "#fff";
                    const text = `Loss: ${currentLoss.toFixed(4)}`;
                    ctx.fillText(text, currentX + 15, legendY);
                    const itemWidth = ctx.measureText(text).width + 20;
                    this.legendHitboxes["chart"] = {
                        x: currentX,
                        y: legendY - 10,
                        width: itemWidth,
                        height: 20
                    };
                    currentX += itemWidth + 15;
                    
                    // 只显示有效数据点的数量
                    ctx.fillStyle = "#fff";
                    ctx.fillText(`数据点: ${validLossData.length}`, currentX, legendY);
                } else {
                    // 如果没有有效的loss数据，不显示任何图例信息
                    // 这样可以避免显示"数据点: 0"或无效的loss值
                }
                
                ctx.globalAlpha = 1.0;
            };
        }
    });

    // 自动注册事件监听，自动分发到updateProgress
    if (eventName && comfyClass) {
        app.api.addEventListener(eventName, (event) => {
            const data = event.detail;
            app.graph.findNodesByType(comfyClass).forEach(node => {
                if (String(data.node_id) === String(node.id)) {
                    node.updateProgress({
                        current_step: data.current_step,
                        total_steps: data.total_steps,
                        status: data.status,
                        loss: data.loss
                    });
                }
            });
        });
    }
} 